import os
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import scipy.io as scio
import torch
import torch.nn as nn
import cv2
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
from torch.utils.data import SubsetRandomSampler
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import cohen_kappa_score
from func import load,product,preprocess
from network import spec_cnn, operate
import time

## 指定显卡

USE_GPU=True
if USE_GPU:
    os.environ["CUDA_VISIBLE_DEVICES"] = "3"
else:
    device=torch.device('cpu')

print('using device:','0')

####################################load dataset(indian_pines & pavia_univ & ksc)######################

a=load()

All_data,labeled_data,rows_num,categories,r,c,flag=a.load_data(flag='ksc')


print('Data has been loaded successfully!')

#################################### 归一化 ######################
#设置归一化范围
mi=-1
ma=1
a=product(flag)
Alldata_norm=a.normlization(All_data[:,1:-1],mi,ma)

#################################### 空间数据，训练、检验、测试、预测 ###########################

#生成训练样本
Experiment_result=np.zeros([categories+4,12])#OA,AA,kappa，重复10次实验

#实验次数
Experiment_num=10

for count in range(0,Experiment_num):

    a=product(flag)
    rows_num,trn_num,tes_num,pre_num=a.generation_num(labeled_data,rows_num,All_data)

    print('第{}次实验，Spatial dataset preparation Finished!'.format(count))

    ################################### numpy2tensor  ####################################

    #label
    y_trn=All_data[trn_num,-1]
    y_tes=All_data[tes_num,-1]
    y_pre=All_data[pre_num,-1]#include background

    print('第{}次实验，Label preparation Finished!'.format(count))

    trn_XX=torch.from_numpy(Alldata_norm[trn_num,:])#

    tes_XX=torch.from_numpy(Alldata_norm[tes_num,:])

    pre_XX=torch.from_numpy(Alldata_norm[pre_num,:])

    trn_YY=torch.from_numpy(y_trn-1)#考虑NLLLoss,标记从0开始
    tes_YY=torch.from_numpy(y_tes-1)
    pre_YY=torch.from_numpy(np.ones([y_pre.shape[0]]))

    print('第{}次实验，Tensor preparation Finished!'.format(count))

    #################################### Training #####################################

    torch.cuda.empty_cache()#GPU memory released

    trn_dataset=TensorDataset(trn_XX,trn_YY)
    trn_loader=DataLoader(trn_dataset,batch_size=20,sampler=SubsetRandomSampler(range(trn_XX.shape[0])))

    #config lr & epoch

    if flag=='indian':
        lr=5e-3
        epoch=700
    elif flag=='pavia':
        lr=1e-2
        epoch=300
    elif flag=='ksc':
        lr=1e-3
        epoch=600
    else:
        raise NotImplementedError

    net = spec_cnn(flag, trn_XX.shape[-1], categories-1, init_weights=True)
    if torch.cuda.device_count() > 1:
        net = nn.DataParallel(net)
    net = net.cuda()
    criterion = torch.nn.NLLLoss()  # 负对数似然损失函数（如果不算log_softmax则直接采用交叉熵损失函数)
    optimizer = torch.optim.SGD(net.parameters(), lr=lr, weight_decay=1e-3)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[202], gamma=0.5)
    loss_trn = []

    trn_time1 = time.time()

    for i in range(1, epoch):
        a=operate()
        loss_trn = a.train(i, loss_trn, net, optimizer, scheduler, trn_loader, criterion)

    trn_time2 = time.time()

    #print(type(loss_trn))  ######CPU
    # plt.figure(1)
    # plt.plot(np.array(loss_trn), label='Training')
    # plt.legend()
    # plt.show()

    ##save training model
    torch.save(net,'1dcnn.pkl')

    print('第{}次实验，模型训练阶段完成！！'.format(count))

    ######################################### inference ####################################

    #y_val=val_YY.numpy()+1#label从1开始
    y_tes=tes_YY.numpy()+1
    y_pre=pre_YY.numpy()+1

    ################### 推断，测试集 ################

    tes_dataset=TensorDataset(tes_XX,tes_YY)
    tes_loader=DataLoader(tes_dataset,batch_size=500)

    #net=torch.load('1dcnn.pkl',map_location='cpu')
    net=net.cuda()
    a=operate()

    tes_time1 = time.time()
    y_pred_tes=a.inference(net,tes_loader,criterion,FLAG='TEST')
    tes_time2 = time.time()

    ####################################### Assess, 测试集 ###########################################

    print('==================Test set=====================')
    print('第{}次实验，测试集OA={}'.format(count,np.mean(y_tes==y_pred_tes)))
    print('第{}次实验，测试集Kappa={}'.format(count,cohen_kappa_score(y_tes,y_pred_tes)))

    ########## 各类别精度

    num_tes=np.zeros([categories-1])
    num_tes_pred=np.zeros([categories-1])
    for k in y_tes:
        num_tes[k-1]=num_tes[k-1]+1
    for j in range(y_tes.shape[0]):
        if y_tes[j]==y_pred_tes[j]:
            num_tes_pred[y_tes[j]-1]=num_tes_pred[y_tes[j]-1]+1

    Acc=num_tes_pred/num_tes*100

    Experiment_result[0,count]=np.mean(y_tes==y_pred_tes)*100#OA
    Experiment_result[1,count]=np.mean(Acc)#AA
    Experiment_result[2,count]=cohen_kappa_score(y_tes,y_pred_tes)*100#Kappa
    Experiment_result[3, count] = trn_time2 - trn_time1
    Experiment_result[4, count] = tes_time2 - tes_time1
    Experiment_result[5:,count]=Acc

    print('第{}次实验，模型评估阶段完成！！!'.format(count))

########## 计算多次实验的均值与标准差并保存

Experiment_result[:,-2]=np.mean(Experiment_result[:,0:-2],axis=1)
Experiment_result[:,-1]=np.std(Experiment_result[:,0:-2],axis=1)

scio.savemat('1dcnn_result_'+str(flag)+'.mat',{'data':Experiment_result})

####################################### Classification map #######################################

######### 推断，预测集 #########

pre_dataset=TensorDataset(pre_XX,pre_YY)
pre_loader=DataLoader(pre_dataset,batch_size=500)

#net=torch.load('1dcnn.pkl',map_location='cpu')
#net=net.module#if use DataParallel

net=net.cuda()
#net=nn.DataParallel(net,device_ids=[0])
#net=net.cpu()
a=operate()
y_pred_pre=a.inference(net,pre_loader,criterion,FLAG='PRED')

print('模型预测阶段完成！！！'.format(count))

#############  展示 ###########

y_disp=np.zeros([All_data.shape[0]])
y_disp[trn_num]=y_trn
y_disp[tes_num]=y_pred_tes

y_disp_gt=y_disp.copy()
y_disp_gt[tes_num]=y_tes

y_disp_all=y_disp.copy()
y_disp_all[pre_num]=y_pred_pre

# plt.subplots(figsize=[10,10])
# ax1=plt.subplot(1,3,1)
# plt.xlabel('TEST')
# a1=plt.imshow(y_disp.reshape(r,c),cmap='jet')
# plt.xticks([])
# plt.yticks([])
#
#
# ax2=plt.subplot(1,3,2)
# plt.xlabel('gt')
# a2=plt.imshow(y_disp_gt.reshape(r,c),cmap='jet')
# plt.xticks([])
# plt.yticks([])
#
#
# ax3=plt.subplot(1,3,3)
# plt.xlabel('All image')
# a3=plt.imshow(y_disp_all.reshape(r,c),cmap='jet')
# plt.xticks([])
# plt.yticks([])
#
# plt.show()

#####################  保存预测全图 #####################

#保存生成的预测图便于后期上色

cv2.imwrite('1dcnn_all_'+str(flag)+'.png', y_disp_all.reshape(r,c))

# plt.subplots(figsize=[10,10])
# a1=plt.imshow(y_disp_all.reshape(r,c),cmap='jet')
# plt.xticks([])
# plt.yticks([])
# plt.savefig('2dcnn_all_'+str(flag)+'.png',dpi=600,bbox_inches='tight')
# plt.show()

print('结果图展示 & 保存阶段完成！！')